#!/usr/bin/env python
# Usage: ./plot.py <candidate_stats_file_path>
import csv
import matplotlib.pyplot as plt
import os
import sys


stats_file = sys.argv[1] 
sizes_post_acceptance =[]
est_gains_accepted = []
act_gains_accepted = []
est_gains_rejected = []
act_gains_rejected = []


def size_plot(sizes):
	plt.xlabel("Number of accepted itemsets")
	plt.ylabel("Relative compression (L%)")
	size_st = sizes[0]
	sizes = [size * 100/size_st for size in sizes]
	plt.plot(sizes)
	plt.savefig("size_vs_is.pdf")

def gain_plot(est_gains_accepted, act_gains_accepted, est_gains_rejected, act_gains_rejected):
	plt.xlabel("Actual Gain (in bits)")
	plt.ylabel("Estimated Gain (in bits)")

	plt.scatter(act_gains_accepted, est_gains_accepted, s=20, c=u'r', marker=u'o', label="accepted candidate")
	plt.scatter(act_gains_rejected, est_gains_rejected, s=20, c=u'b', marker=u'x', label="rejected candidate")
	plt.legend()
	plt.savefig("act_vs_est_gain.pdf")

with open(stats_file, 'r') as fp:
	reader = csv.reader(fp)
	for row in reader:
		if reader.line_num == 1:
			continue
		accepted = row[3] == '1'
		est_gain = float(row[1])
		act_gain = float(row[2])
		size = float(row[4])

		if accepted:
			sizes_post_acceptance.append(size)
			act_gains_accepted.append(act_gain)
			est_gains_accepted.append(est_gain)
		else:
			act_gains_rejected.append(act_gain)
			est_gains_rejected.append(est_gain)

if len(sizes_post_acceptance) > 0:
    size_plot(sizes_post_acceptance)
    gain_plot(est_gains_accepted, act_gains_accepted, est_gains_rejected, act_gains_rejected)
